package com.ai.aicodeinghelper.controller;

import com.ai.aicodeinghelper.advisors.MyLoggerAdvisor;
import com.ai.aicodeinghelper.tools.TimeTools;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.alibaba.cloud.ai.memory.jdbc.MysqlChatMemoryRepository;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.*;

/**
 * 实现了4中方式
 * 1. 没有流式输出
 * 2. Flux<String> 流式输出
 * 3. Flux + ServerSendEvent 流式输出
 * 4. SseEmitter 流式输出
 */
@RestController
@RequestMapping("/helloword")
@Slf4j
public class ChatController {

    private static final String DEFAULT_PROMPT = "你是一个博学的智能聊天助手，请根据用户提问回答";
    //    private static final String DEFAULT_PROMPT = "你是一个友好的聊天机器人，用{voice}的声音回答问题";
    private final ChatClient dashScopeChatClient;

    private final ChatClient ollamaChatClient;

    private final ChatClient openAIChatClient;

    public ChatController(JdbcTemplate jdbcTemplate, DashScopeChatModel chatModel, OllamaChatModel ollamaChatModel,ToolCallback[] allTools,ChatModel openAiChatModel) {
        System.out.println("初始化");

        // 构造 ChatMemoryRepository 和 ChatMemory
        // 基于Mysql的记忆
        ChatMemoryRepository chatMemoryRepository = MysqlChatMemoryRepository.mysqlBuilder()
                .jdbcTemplate(jdbcTemplate)
                .build();
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .build();


        // 基于内存的记忆
        InMemoryChatMemoryRepository inMemoryChatMemoryRepository = new InMemoryChatMemoryRepository();
        MessageWindowChatMemory chatMemoryInMemroy = MessageWindowChatMemory.builder().chatMemoryRepository(inMemoryChatMemoryRepository).maxMessages(20).build();

        this.dashScopeChatClient = ChatClient.builder(chatModel)
                .defaultSystem(DEFAULT_PROMPT)
                .defaultOptions(DashScopeChatOptions
                        .builder()
                        .withTopP(0.7)
                        .build())
                .defaultToolCallbacks(allTools)
                .defaultAdvisors(new MyLoggerAdvisor())
//                .defaultAdvisors(new SimpleLoggerAdvisor())
//                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build())
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemoryInMemroy).build())
                .build();
        this.ollamaChatClient = ChatClient.builder(ollamaChatModel)
                .defaultSystem(DEFAULT_PROMPT)
                .defaultToolCallbacks(allTools)
                .defaultOptions(ChatOptions.builder().topP(0.7).build())
                .defaultAdvisors(new MyLoggerAdvisor(), MessageChatMemoryAdvisor.builder(chatMemoryInMemroy).build())
                .build();
        this.openAIChatClient = ChatClient.builder(openAiChatModel)
                .defaultSystem(DEFAULT_PROMPT)
                .defaultToolCallbacks(allTools)
                .defaultOptions(ChatOptions.builder().topP(0.7).build())
                .defaultAdvisors(new MyLoggerAdvisor(), MessageChatMemoryAdvisor.builder(chatMemoryInMemroy).build())
                .build();
    }

    @GetMapping("/simple/chat")
    public String simpleChat(@RequestParam(name = "query", defaultValue = "你好，很高兴认识你，能简单介绍下自己吗？") String query,
                             @RequestParam(name = "chat-id", defaultValue = "1") String chatId) {
        log.info("testPrompt =>{}", query);
        ChatResponse chatResponse = openAIChatClient.prompt().user(query).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId)).call().chatResponse();
        String content = chatResponse.getResult().getOutput().getText();
        ChatResponseMetadata metadata1 = chatResponse.getMetadata();
        List<Generation> results = chatResponse.getResults();
        log.info("prompt content =>{}", content);
        log.info("metadata =>{}", metadata1);
        log.info("results =>{}", results.size());

        return content;
    }

    //    @GetMapping(value = "/stream/chat",produces = {"text/html;charset=UTF-8"})
    @GetMapping(value = "/stream/chat")
    public Flux<String> streamChat(@RequestParam(name = "query", defaultValue = "你好，很高兴认识你，能简单介绍下自己吗？") String query,
                                   @RequestParam(name = "chat-id", defaultValue = "1") String chatId,
                                   HttpServletResponse response) {
        response.setCharacterEncoding("UTF-8"); // 处理中文乱码两种方式都可
        Flux<String> content = dashScopeChatClient.prompt()
                .system(a -> a.param("voice", "小团团"))
                .user(query).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId)).stream().content();
        return content;
    }

    @GetMapping("/sse/chat")
    public Flux<ServerSentEvent<String>> sseChat(@RequestParam(name = "query", defaultValue = "你好，很高兴认识你，能简单介绍下自己吗？") String query,
                                                 @RequestParam(name = "chat-id", defaultValue = "1") String chatId,
                                                 HttpServletResponse response) {
        response.setCharacterEncoding("UTF-8");
        Flux<String> content = dashScopeChatClient.prompt().user(query).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId)).stream().content();
        Flux<ServerSentEvent<String>> map = content.map(chunk -> {
            return ServerSentEvent.<String>builder()
                    .data(chunk) // 如果需要额外信息这里可以传递一个JSON串
                    .build();
        });
        return map;
    }

    @GetMapping("/sse_emitter/chat")
    public SseEmitter sseEmitterChat(@RequestParam(name = "query", defaultValue = "你好，很高兴认识你，能简单介绍下自己吗？") String query,
                                     @RequestParam(name = "chat-id", defaultValue = "1") String chatId,
                                     HttpServletResponse response) {
        response.setCharacterEncoding("UTF-8");
        SseEmitter sseEmitter = new SseEmitter();
        Flux<String> content = dashScopeChatClient.prompt().user(query).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId)).stream().content();
        content.subscribe(chunk -> {
            try {
                sseEmitter.send(chunk);
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
                throw new RuntimeException(e);
            }
        }, sseEmitter::completeWithError, sseEmitter::complete);
        return sseEmitter;
    }

    record ChatReport(String title, List<String> list) {

    }

    @GetMapping("/chat/report")
    public ChatReport chatReport(@RequestParam(name = "chat-id", defaultValue = "1") String chatId) {
        ChatReport entity = dashScopeChatClient.prompt()
                .system(DEFAULT_PROMPT + "每次对话后都要生成结果，标题为{title}的报告，内容建议为列表")
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId)).call()
                .entity(ChatReport.class);
        log.info("report entity =>{}", entity);
        return entity;
    }

    record ActorFilms(String actor,List<String> movies){}

    // 获取刘亦菲演过的电影
    @GetMapping("/movie")
    public ActorFilms movie(){
        return dashScopeChatClient.prompt().user("刘亦菲").call().entity(ActorFilms.class);
    }

    @GetMapping("/movies")
    public List<ActorFilms> movies(){
        return dashScopeChatClient.prompt().user("生成5部刘亦菲的电影集").call().entity(new ParameterizedTypeReference<List<ActorFilms>>() {});
    }

    /**
     * 动态切换某个平台的多个模型
     *
     * @param model
     * @param query
     * @param chatId
     * @return
     */
    @GetMapping("/model/chat")
    public Flux<String> modelChat(@RequestParam(name = "model", defaultValue = "qwen-plus") String model,
                                  @RequestParam(name = "query") String query,
                                  @RequestParam(name = "chat-id", defaultValue = "1") String chatId) {
        // 也可以使用这种方式定义Options
//        ChatOptions options = ChatOptions.builder().model(model).build();

        DashScopeChatOptions dashScopeChatOptions = DashScopeChatOptions.builder().withModel(model).withTopP(0.7).build();

        log.info("使用的模型 =>{}", dashScopeChatOptions.getModel());
        Flux<String> content = this.dashScopeChatClient.prompt()
                // options 会替代defaultOptions
                // 并且如果有多个Options会被覆盖，只会生效最后一个
                .options(dashScopeChatOptions)
                .user(query)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .stream()
                .content();
        return content;
    }


    /**
     * 动态切换平台
     *
     * @param platform
     * @param query
     * @return
     */
    @GetMapping("/platform/chat")
    public Flux<String> platformChat(@RequestParam(name = "platform", defaultValue = "dashScope") String platform,
                                     @RequestParam(name = "query") String query,
                                     @RequestParam(name = "chat-id", defaultValue = "1") String chatId
    ) {
        if ("dashScope".equals(platform)) {
            return this.dashScopeChatClient.prompt().user(query).stream().content();
        }

        if ("ollama".equals(platform)) {
            Flux<ChatResponse> chatResponseFlux = this.ollamaChatClient.prompt()
                    .tools(new TimeTools())
                    .user(query)
                    .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                    .stream().chatResponse();
            return chatResponseFlux
                    .map(a -> a.getResult().getOutput().getText());
//                    .concatWith(Flux.just("最终状态标识"));
        }

        throw new RuntimeException("暂无平台信息");
    }

}
